import numpy as np
from scipy import misc
import tensorflow as tf
import cv2, re, base64

# Read image from base64 string
def read_img_base64(base64str):
	# check if contains javascript
	image_data = re.sub('^data:image/.+;base64,', '', base64str).decode('base64')
	arr = np.fromstring(image_data, np.uint8)
	img = cv2.imdecode(arr, cv2.IMREAD_COLOR).astype(np.float32)
	return img

# Read image from base64 string
def read_img_blob(im_str):
	# check if contains javascript
	arr = np.fromstring(im_str, dtype=np.uint8)
	img = cv2.imdecode(arr, cv2.IMREAD_COLOR).astype(np.float32)
	return img

# VGG 16 accepts RGB channel 0 to 1 (This tensorflow model).
def load_image_array(image_file):
	img = misc.imread(image_file)
	return get_image_array(img)

def get_image_array(img):
	# GRAYSCALE
	if len(img.shape) == 2:
		img_new = np.ndarray( (img.shape[0], img.shape[1], 3), dtype = 'float32')
		img_new[:,:,0] = img
		img_new[:,:,1] = img
		img_new[:,:,2] = img
		img = img_new
	img_resized = misc.imresize(img, (224, 224))
	return (img_resized/255.0).astype('float32')


# FOR PREDICTION ON A SINGLE IMAGE
def get_image_feed(img):
	image_array = get_image_array(img)
	image_feed = np.ndarray((1,224,224,3))
	image_feed[0:,:,:] = image_array

	return image_feed

def get_vgg_graph(model_path):
	with open(model_path) as vgg_file:
		vgg16raw = vgg_file.read()

	graph_def = tf.GraphDef()
	graph_def.ParseFromString(vgg16raw)
	images = tf.placeholder("float32", [None, 224, 224, 3])
	tf.import_graph_def(graph_def, input_map={ "images": images })
	graph = tf.get_default_graph()
	fc7_tensor = graph.get_tensor_by_name("import/Relu_1:0")
	tf.reset_default_graph()

	return graph, fc7_tensor, images

def get_fc7_feas(img, vgg_graph, fc7_tensor, images):
	image_feed = get_image_feed(img)
	with tf.Session(graph = vgg_graph) as sess:
		feed_dict  = { images : image_feed }
		fc7_features = sess.run(fc7_tensor, feed_dict = feed_dict)

	return fc7_features


def extract_fc7_features(image_path, model_path):
	vgg_file = open(model_path)
	vgg16raw = vgg_file.read()
	vgg_file.close()

	graph_def = tf.GraphDef()
	graph_def.ParseFromString(vgg16raw)
	images = tf.placeholder("float32", [None, 224, 224, 3])
	tf.import_graph_def(graph_def, input_map={ "images": images })
	graph = tf.get_default_graph()

	# sess = tf.Session()
	with tf.Session() as sess:
		image_array = load_image_array(image_path)
		image_feed = np.ndarray((1,224,224,3))
		image_feed[0:,:,:] = image_array
		feed_dict  = { images : image_feed }
		fc7_tensor = graph.get_tensor_by_name("import/Relu_1:0")
		fc7_features = sess.run(fc7_tensor, feed_dict = feed_dict)
	# sess.close()

	return fc7_features